# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     test
   Description :   
   Author :       lth
   date：          2022/1/27
-------------------------------------------------
   Change Activity:
                   2022/1/27 13:35: create this script
-------------------------------------------------
this script is used for the test, its main purpose is to convenient the developer test
and everyone can use this model easily
"""
__author__ = 'lth'

import cv2
import numpy as np
import torch
from torch.backends import cudnn

from config import GetConfig
from datalist import test_transform, UnetTableData
from model2 import UNet
from utils import PostProcess

metric_loss = 100
scale = 800
PostProcess = PostProcess()


class ResizeSuitable:

    def __init__(self):
        ResizeSuitable.height_width_min_value = 1000

    @staticmethod
    def resize(data):
        im = data
        h, w, _ = im.shape
        # if min(h, w) > ResizeSuitable.height_width_min_value:
        #     ratio = h / w
        #     if h > w:
        #         h = ResizeSuitable.height_width_min_value
        #         w = int(h / ratio)
        #     else:
        #         w = ResizeSuitable.height_width_min_value
        #         h = int(h * ratio)
        h = h // 64 * 32
        w = w // 64 * 32
        im = cv2.resize(im, (w, h))
        # im = cv2.copyMakeBorder(im,(0,0,0,0),cv2.BORDER_CONSTANT,(255,255,255))
        data = im
        return data


class UNetTable:
    def __init__(self):
        self.args = GetConfig()

        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")

        '''
        定义模型
        '''

        self.model = UNet(n_channels=3, num_classes=4).to(self.device)

        if use_cuda:
            self.generator = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True
            cudnn.enabled = True

        print("load the weight from pretrained-weight file")
        model_dict = self.generator.state_dict()
        checkpoint = torch.load(self.args.pretrained_weight)['model_state_dict']
        model_dict.update(checkpoint)
        self.generator.load_state_dict(model_dict, strict=True)
        print("Restoring the weight from pretrained-weight file \nFinished loading the weight")

        self.generator.eval()
        self.resizer = ResizeSuitable()

    @torch.no_grad()
    def inference(self, image_path):
        img = cv2.imread(image_path)
        img = self.resizer.resize(img)
        temp = img
        img = test_transform(img)
        img = img.unsqueeze(0).to(self.device)
        result = self.generator(img)

        return result, temp


if __name__ == "__main__":
    model = UNetTable()
    import os

    for image in os.listdir("test_data"):

        output, temp = model.inference("test_data/" + image)

        image = UnetTableData.decode_label_map(output[0], threshold=0.5)

        target = np.zeros_like(image[0])

        for index, i in enumerate(image):
            if index == 0 or index == 1:
                v = True
            else:
                v = False
            i = PostProcess.process(i, v=v, index=index)
            target = i | target

        # target=cv2.cvtColor(target,cv2.COLOR_RGB2BGR)
        cv2.addWeighted(target, 0.5, temp, 0.5, 0, target)
        cv2.imshow("Dsa", target)
        cv2.waitKey()
        cv2.imwrite("result.jpg", target)
